import os.path

import tensorboard
from matplotlib import pyplot as plt
from tensorflow.keras.layers import Dense, Convolution2D, MaxPooling2D, Dropout, Activation, Flatten, BatchNormalization
from tensorflow.keras.models import Sequential
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.optimizers import Adam, RMSprop
from tensorflow.keras.losses import CategoricalCrossentropy
from tensorflow.keras.metrics import Accuracy, Precision, Recall
from tensorflow.keras.callbacks import ModelCheckpoint

def build_model():
    model = Sequential()

    model.add(Convolution2D(
        filters=32,
        kernel_size=(3,3),
        strides=(1,1),
        padding='same',
        input_shape=(150,150,3)
    ))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Convolution2D(
        filters=32,
        kernel_size=(3, 3),
        strides=(1, 1),
        padding='same'
    ))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(MaxPooling2D(
        pool_size=(2,2),
        strides=(2,2),
        padding='same'
    ))
    model.add(Dropout(0.1))

    model.add(Convolution2D(
        filters=64,
        kernel_size=(3, 3),
        strides=(1, 1),
        padding='same'
    ))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Convolution2D(
        filters=64,
        kernel_size=(3, 3),
        strides=(1, 1),
        padding='same'
    ))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(MaxPooling2D(
        pool_size=(2, 2),
        strides=(2, 2),
        padding='same'
    ))
    model.add(Dropout(0.1))

    model.add(Convolution2D(
        filters=128,
        kernel_size=(3, 3),
        strides=(1, 1),
        padding='same'
    ))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Convolution2D(
        filters=128,
        kernel_size=(3, 3),
        strides=(1, 1),
        padding='same'
    ))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(MaxPooling2D(
        pool_size=(2, 2),
        strides=(2, 2),
        padding='same'
    ))
    model.add(Dropout(0.1))

    model.add(Flatten())

    model.add(Dense(512))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Dropout(0.5))

    model.add(Dense(128))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Dropout(0.5))

    model.add(Dense(2))
    model.add(BatchNormalization())
    model.add(Activation('softmax'))

    model.summary()
    return model
def buildModel():
    conModel = Sequential()
    # 定义卷积层 第一层要加importShape 不加会报错
    conModel.add(Convolution2D(32, (3, 3), (1, 1), padding='same', input_shape=(150, 150, 3)))
    conModel.add(Convolution2D(32, (3, 3), (1, 1),padding='same'))
    # 下采样
    conModel.add(MaxPooling2D((3, 3), (2, 2), padding='valid'))
    conModel.add(Activation('relu'))

    conModel.add(Convolution2D(64, (3, 3), (1, 1), padding='same'))
    conModel.add(Convolution2D(64, (3, 3), (1, 1),padding='same'))
    # 下采样
    conModel.add(MaxPooling2D((3, 3), (2, 2), padding='valid'))
    conModel.add(Activation('relu'))

    conModel.add(Convolution2D(128, (3, 3), (1, 1),padding='same'))
    conModel.add(Convolution2D(128, (3, 3), (1, 1),padding='same'))
    # 下采样
    conModel.add(MaxPooling2D((3, 3), (2, 2), padding='valid'))
    conModel.add(Activation('relu'))

    conModel.add(Convolution2D(256, (3, 3), (1, 1), padding='same'))
    conModel.add(Convolution2D(256, (3, 3), (1, 1), padding='same'))
    # 下采样
    conModel.add(MaxPooling2D((3, 3), (2, 2), padding='valid'))
    conModel.add(Activation('relu'))
    # 输出结构
    #     后边其实还可以加多几层， 自已可以去实验
    #     定义完前边的卷积层就能提取图像的信息， 后边要加上全连接
    fc = Sequential()
    #     对应卷积层的输出
    #     展平向量
    fc.add(Flatten())
    fc.add(Dense(256))
    fc.add(BatchNormalization())
    fc.add(Activation('relu'))
    #     失活防止过拟合
    fc.add(Dropout(0.4))
    fc.add(Dense(128))
    fc.add(BatchNormalization())
    fc.add(Activation('relu'))
    #     失活防止过拟合
    fc.add(Dropout(0.4))
    #     对应的图像种类
    fc.add(Dense(2))
    fc.add(Activation('softmax'))

    model = Sequential()
    model.add(conModel)
    model.add(fc)

    model.summary()
    return model


def getTrainData():
    # 这里拿之前划分好的测试集训练不然图像太多了 train大概25000 val 2500
    train_root = os.path.join('../../train/val')
    # 拿keras的图像迭代器获取数据集和标签， 你也可以自已写一个
    # 归一化 验证集划分0.1
    train_genera = ImageDataGenerator(rescale=1./255, validation_split=0.1)
    # class_mode代表分类 会返回标签001 010 100 这样的 subset 标注为训练集， 刚刚上边划分的validation_split
    train_gen = train_genera.flow_from_directory(directory=train_root,
                                                 target_size=(150, 150),
                                                 class_mode='categorical',
                                                 subset='training',
                                                 batch_size=64)
    val_gen = train_genera.flow_from_directory(directory=train_root,
                                               target_size=(150, 150),
                                               class_mode='categorical',
                                               subset='validation',
                                               batch_size=64)
    return train_gen, val_gen


if __name__ == "__main__":
    model = build_model()
    # 模型指标
    # 无脑Adam
    # Adam优化器 学习率 0.0001
    # CategoricalCrossentropy损失函数
    # metrics 训练时候输出的加就完事了
    model.compile(optimizer=RMSprop(learning_rate=0.0005)
                  , loss=CategoricalCrossentropy(from_logits=False)
                  , metrics=[Accuracy(), Recall(), Precision()])
    # 回调函数 val_loss 评估指标 要有测试集才能使
    cb = ModelCheckpoint('dogCat.h5',
                         monitor='accuracy',
                         save_best_only=True,
                         save_weights_only=True)
    # 暂时忘了
    # tenCb = tensorboard.
    train_gen, val_gen = getTrainData()
    print(train_gen.class_indices)
    print(val_gen.class_indices)
    # 模型训练 train_gen 包含 图像和标签
    res = model.fit(x=train_gen,
                    steps_per_epoch=train_gen.n // 64,
                    batch_size=64,
                    epochs=10,
                    callbacks=cb
                    )
    print(res)
    res = res.history
    for i in res:
        print(i)
        print(res[i])
        plt.figure()
        plt.plot(range(10), res[i],)
        plt.ylabel(i)
        plt.xlabel('epochs')
        plt.savefig(i + '.jpg')
